import os
import time
import uuid

import numpy as np

import cv2
from django.conf import settings

from apps.img.models import SplicingInfo, SplicingImgInfo, imgInfo


def splicingRow(img1, img2, mainImg):
    MIN = 10
    surf = cv2.xfeatures2d.SIFT_create(400)

    kp1, descrip1 = surf.detectAndCompute(img1, None)
    kp2, descrip2 = surf.detectAndCompute(img2, None)

    FLANN_INDEX_KDTREE = 0
    indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    searchParams = dict(checks=50)

    flann = cv2.FlannBasedMatcher(indexParams, searchParams)
    match = flann.knnMatch(descrip1, descrip2, k=2)

    good = []
    for i, (m, n) in enumerate(match):
        if m.distance < n.distence:
            good.append(m)

    if len(good) > MIN:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        ano_pts = np.float32([kp2[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        M, mask = cv2.findHomography(src_pts, ano_pts, cv2.RANSAC, 5.0)
        warpImg = cv2.warpPerspective(img2, np.linalg.inv(M), (img1.shape[1] + img2.shape[1], img2.shape[0]))
        # direct = warpImg.copy()
        # direct[0:img1.shape[0], 0:img1.shape[1]] = img1
        simple = time.time()

        rows, cols = img1.shape[:2]

        for col in range(0, cols):
            if img1[:, col].any() and warpImg[:, col].any():  # 开始重叠的最左端
                left = col
                break
        for col in range(cols - 1, 0, -1):
            if img1[:, col].any() and warpImg[:, col].any():  # 重叠的最右一列
                right = col
                break

        res = np.zeros([rows, cols, 3], np.uint8)
        for row in range(0, rows):
            for col in range(0, cols):
                if not img1[row, col].any():  # 如果没有原图，用旋转的填充
                    res[row, col] = warpImg[row, col]
                elif not warpImg[row, col].any():
                    res[row, col] = img1[row, col]
                else:
                    srcImgLen = float(abs(col - left))
                    testImgLen = float(abs(col - right))
                    alpha = srcImgLen / (srcImgLen + testImgLen)
                    res[row, col] = np.clip(img1[row, col] * (1 - alpha) + warpImg[row, col] * alpha, 0, 255)

        warpImg[0:mainImg.shape[0], 0:mainImg.shape[1]] = res
        final = time.time()
        img3 = cv2.cvtColor(warpImg, cv2.COLOR_BGR2RGB)
        return img3
    else:
        return None


def splicingCol(img1, img2, mainImg):
    MIN = 10
    surf = cv2.xfeatures2d.SIFT_create(400)  # 可以改为SIFT

    kp1, descrip1 = surf.detectAndCompute(img1, None)
    kp2, descrip2 = surf.detectAndCompute(img2, None)

    FLANN_INDEX_KDTREE = 0
    indexParams = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    searchParams = dict(checks=50)

    flann = cv2.FlannBasedMatcher(indexParams, searchParams)
    match = flann.knnMatch(descrip1, descrip2, k=2)

    good = []
    for i, (m, n) in enumerate(match):
        if m.distance < 0.75 * n.distance:
            good.append(m)

    if len(good) > MIN:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        ano_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
        M, mask = cv2.findHomography(src_pts, ano_pts, cv2.RANSAC, 5.0)
        warpImg = cv2.warpPerspective(img2, np.linalg.inv(M), (img2.shape[1], img1.shape[0] + img2.shape[0]))

        rows, cols = img1.shape[:2]

        for row in range(0, rows):
            if img1[row, :].any() and warpImg[row, :].any():  # 开始重叠的最左端
                top = row
                break
        for row in range(rows - 1, 0, -1):
            if img1[row, :].any() and warpImg[row, :].any():  # 重叠的最右一列
                button = row
                break

        res = np.zeros([rows, cols, 3], np.uint8)
        for col in range(0, cols):
            for row in range(0, rows):
                if not img1[row, col].any():  # 如果没有原图，用旋转的填充
                    res[row, col] = warpImg[row, col]
                elif not warpImg[row, col].any():
                    res[row, col] = img1[row, col]
                else:
                    srcImgLen = float(abs(row - top))
                    testImgLen = float(abs(row - button))
                    alpha = srcImgLen / (srcImgLen + testImgLen)
                    res[row, col] = np.clip(img1[row, col] * (1 - alpha) + warpImg[row, col] * alpha, 0, 255)

        warpImg[0:mainImg.shape[0], 0:mainImg.shape[1]] = res

        img4 = cv2.cvtColor(warpImg, cv2.COLOR_BGR2RGB)

        return img4
    else:
        return None


def splicing(splicingId):
    splicingInfo = SplicingInfo.objects.get(id=splicingId)
    row = int(splicingInfo.splicing_row)
    col = int(splicingInfo.splicing_col)
    totalSplicingCount = row * (col - 1) + row - 1
    count = 0
    imgDatas = []
    for i in range(1, row):
        imgs = SplicingImgInfo.objects.getByRow(sInfoId=splicingInfo, row=row)
        imgList = []
        for data in imgs:
            img = data.img_info
            imgList.append(img.file_path)
        imgDatas.append(imgList)
    print(imgDatas)
    rowList = []
    for rowImg in imgDatas:
        centerCol = int(len(rowImg) / 2)
        leftList = []
        rightList = []
        for i in range(len(rowImg)):
            if i <= centerCol:
                leftList.append(rowImg[i])
            else:
                rightList.append(rowImg[i])
        urlA = leftList[0]
        imgLeft = cv2.imread(urlA)
        for urlB in leftList[1:]:
            imgB = cv2.imread(urlB)
            tmp = splicingRow(imgLeft, imgB, imgB)
            imgLeft = tmp
            count += 1
            process = int((count / totalSplicingCount) * 100)
            splicingInfo.progress = process
            splicingInfo.save()
        for right in rightList:
            imgB = cv2.imread(right)
            tmp = splicingRow(imgLeft, imgB, imgLeft)
            imgLeft = tmp
            count += 1
            process = int((count / totalSplicingCount) * 100)
            splicingInfo.progress = process
            splicingInfo.save()
        rowList.append(imgLeft)

    vCenter = int(len(rowList) / 2)
    topList = []
    bottomList = []

    for i in range(len(rowList)):
        if i <= vCenter:
            topList.append(rowList[i])
        else:
            bottomList.append(rowList[i])
    topImg = cv2.imread(topList[0])
    for urlT in topList[1:]:
        imgT = cv2.imread(urlT)
        tmp = splicingCol(topImg, imgT, imgT)
        topImg = tmp
        count += 1
        process = int((count / totalSplicingCount) * 100)
        splicingInfo.progress = process
        splicingInfo.save()
    for bottom in bottomList:
        imgB = cv2.imread(bottom)
        tmp = splicingCol(topImg, imgB, topImg)
        topImg = tmp
        count += 1
        process = int((count / totalSplicingCount) * 100)
        splicingInfo.progress = process
        splicingInfo.save()
    file_path = os.path.join(settings.BASE_DIR, "media/")
    file_name = splicingInfo.id + '.jpg'
    cv2.imwrite(file_path + file_name, topImg)
    splicingImg = imgInfo()
    splicingImg.id = str(uuid.uuid1())
    splicingImg.file_name = file_name
    splicingImg.file_path = file_path + file_name
    splicingImg.save()

    splicingInfo.splicing_img = splicingImg
    splicingInfo.progress = 100
    splicingInfo.save()
